import os
import pandas as pd
import re

def extract_atom_number(folder_name):
    """
    从文件夹名称中提取原子数量
    例如: atom12000 -> 12000
    """
    match = re.search(r'atom(\d+)', folder_name)
    if match:
        return int(match.group(1))
    return None

def calculate_precision_error_for_folder(folder_path):
    """
    计算指定文件夹下所有CSV文件中comm_type='Total_C-lop_prediction_comm_type_51'的平均precision_error
    """
    precision_errors = []
    
    # 遍历文件夹中的所有CSV文件
    for file in os.listdir(folder_path):
        if file.endswith('.csv'):
            file_path = os.path.join(folder_path, file)
            try:
                # 读取CSV文件
                df = pd.read_csv(file_path)
                
                # 筛选comm_type为'Total_C-lop_prediction_comm_type_51'的行
                target_rows = df[df['comm_type'] == 'Total_C-lop_prediction_comm_type_51']
                
                # 提取precision_error值（确保是数值类型）
                for _, row in target_rows.iterrows():
                    precision_error = pd.to_numeric(row['precision_error'], errors='coerce')
                    if not pd.isna(precision_error):
                        precision_errors.append(precision_error)
                        
            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")
                continue
    
    # 计算平均值
    if precision_errors:
        return sum(precision_errors) / len(precision_errors)
    else:
        return None

def process_diff_atom_accuracy(base_dir):
    """
    处理不同原子数量的精度数据
    """
    results = []
    
    print(f"开始处理目录: {base_dir}")
    
    # 遍历base_dir下的所有子文件夹
    for item in os.listdir(base_dir):
        item_path = os.path.join(base_dir, item)
        
        # 检查是否为目录且符合atomX格式
        if os.path.isdir(item_path) and item.startswith('atom'):
            atom_number = extract_atom_number(item)
            
            if atom_number is not None:
                print(f"处理文件夹: {item} (原子数: {atom_number})")
                
                # 计算该文件夹的平均precision_error
                avg_precision_error = calculate_precision_error_for_folder(item_path)
                
                if avg_precision_error is not None:
                    results.append({
                        'atom_number': atom_number,
                        'error': avg_precision_error
                    })
                    print(f"  - 平均精度误差: {avg_precision_error:.6f}")
                else:
                    print(f"  - 未找到有效的精度数据")
            else:
                print(f"跳过文件夹: {item} (不符合atomX格式)")
    
    # 按原子数量排序
    results.sort(key=lambda x: x['atom_number'])
    
    return results

def save_results(base_dir, results):
    """
    保存结果到CSV文件
    """
    # 创建输出目录
    output_dir = os.path.join(base_dir, 'all_predict_precison')
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建DataFrame并保存
    df = pd.DataFrame(results)
    output_file = os.path.join(output_dir, 'summary.csv')
    df.to_csv(output_file, index=False)
    
    print(f"\n结果已保存到: {output_file}")
    print(f"总共处理了 {len(results)} 个原子配置")
    
    # 显示结果预览
    print("\n结果预览:")
    print(df.to_string(index=False))

if __name__ == '__main__':
    # 设置基础目录
    base_dir = r"F:\PostGraduate\Point-to-Point-Code\App_Prediction\code\ML_Predict\results_with_cache\diff_atom_in_proc32"
    
    try:
        # 处理数据
        results = process_diff_atom_accuracy(base_dir)
        
        if results:
            # 保存结果
            save_results(base_dir, results)
        else:
            print("未找到任何有效数据")
            
    except Exception as e:
        print(f"执行过程中出错: {e}")
